"""
双阈值端点检测

实现了基于能量的双阈值端点检测算法，包括：
- 动态阈值计算
- 中心截幅处理
- 过零率计算
- 中值滤波平滑
"""

import numpy as np
from typing import Tuple, Optional
from loguru import logger
from scipy import signal


class EndpointDetection:
    """双阈值端点检测算法"""
    
    def __init__(self, 
                 energy_threshold_high: float = 0.8,
                 energy_threshold_low: float = 0.3,
                 zero_crossing_threshold: float = 0.5,
                 min_silence_length: float = 0.1,
                 min_speech_length: float = 0.05,
                 center_clip_threshold: float = 0.01,
                 frame_length: int = 256,
                 frame_shift: int = 128,
                 sample_rate: int = 100):
        """
        初始化端点检测器
        
        Args:
            energy_threshold_high: 能量高阈值
            energy_threshold_low: 能量低阈值
            zero_crossing_threshold: 过零率阈值
            min_silence_length: 最小静音长度(秒)
            min_speech_length: 最小语音长度(秒)
            center_clip_threshold: 中心截幅阈值
            frame_length: 帧长度
            frame_shift: 帧移
            sample_rate: 采样率
        """
        self.energy_threshold_high = energy_threshold_high
        self.energy_threshold_low = energy_threshold_low
        self.zero_crossing_threshold = zero_crossing_threshold
        self.min_silence_length = min_silence_length
        self.min_speech_length = min_speech_length
        self.center_clip_threshold = center_clip_threshold
        self.frame_length = frame_length
        self.frame_shift = frame_shift
        self.sample_rate = sample_rate
        
        # 转换为帧数
        self.min_silence_frames = int(min_silence_length * sample_rate / frame_shift)
        self.min_speech_frames = int(min_speech_length * sample_rate / frame_shift)
        
        logger.info("初始化双阈值端点检测器")
    
    def detect_endpoints(self, signal_data: np.ndarray) -> Tuple[int, int]:
        """
        检测信号的端点
        
        Args:
            signal_data: 输入信号
            
        Returns:
            (起始点, 结束点)
        """
        # 预处理信号
        processed_signal = self._preprocess_signal(signal_data)
        
        # 计算短时能量
        energy = self._calculate_short_time_energy(processed_signal)
        
        # 计算过零率
        zero_crossing_rate = self._calculate_zero_crossing_rate(processed_signal)
        
        # 中值滤波平滑
        energy_smooth = self._median_filter(energy)
        zcr_smooth = self._median_filter(zero_crossing_rate)
        
        # 动态阈值计算
        energy_threshold_high, energy_threshold_low = self._calculate_dynamic_thresholds(energy_smooth)
        
        # 端点检测
        start_frame, end_frame = self._detect_boundaries(
            energy_smooth, zcr_smooth, energy_threshold_high, energy_threshold_low
        )
        
        # 转换为采样点
        start_sample = start_frame * self.frame_shift
        end_sample = min(end_frame * self.frame_shift + self.frame_length, len(signal_data))
        
        logger.info(f"检测到端点: {start_sample} - {end_sample}")
        return start_sample, end_sample
    
    def _preprocess_signal(self, signal_data: np.ndarray) -> np.ndarray:
        """
        预处理信号
        
        Args:
            signal_data: 原始信号
            
        Returns:
            预处理后的信号
        """
        # 去除直流分量
        signal_centered = signal_data - np.mean(signal_data)
        
        # 幅值归一化
        max_amplitude = np.max(np.abs(signal_centered))
        if max_amplitude > 0:
            signal_normalized = signal_centered / max_amplitude
        else:
            signal_normalized = signal_centered
        
        return signal_normalized
    
    def _calculate_short_time_energy(self, signal_data: np.ndarray) -> np.ndarray:
        """
        计算短时能量
        
        Args:
            signal_data: 输入信号
            
        Returns:
            短时能量序列
        """
        num_frames = (len(signal_data) - self.frame_length) // self.frame_shift + 1
        energy = np.zeros(num_frames)
        
        for i in range(num_frames):
            start = i * self.frame_shift
            end = start + self.frame_length
            frame = signal_data[start:end]
            energy[i] = np.sum(frame ** 2)
        
        return energy
    
    def _calculate_zero_crossing_rate(self, signal_data: np.ndarray) -> np.ndarray:
        """
        计算过零率
        
        Args:
            signal_data: 输入信号
            
        Returns:
            过零率序列
        """
        # 中心截幅处理
        signal_clipped = self._center_clip(signal_data)
        
        num_frames = (len(signal_clipped) - self.frame_length) // self.frame_shift + 1
        zcr = np.zeros(num_frames)
        
        for i in range(num_frames):
            start = i * self.frame_shift
            end = start + self.frame_length
            frame = signal_clipped[start:end]
            
            # 计算过零次数
            signs = np.sign(frame)
            zero_crossings = np.sum(np.abs(signs[1:] - signs[:-1])) / 2
            zcr[i] = zero_crossings / len(frame)
        
        return zcr
    
    def _center_clip(self, signal_data: np.ndarray) -> np.ndarray:
        """
        中心截幅处理
        
        Args:
            signal_data: 输入信号
            
        Returns:
            截幅后的信号
        """
        clipped = np.copy(signal_data)
        mask = np.abs(signal_data) < self.center_clip_threshold
        clipped[mask] = 0
        return clipped
    
    def _median_filter(self, data: np.ndarray, window_size: int = 5) -> np.ndarray:
        """
        中值滤波
        
        Args:
            data: 输入数据
            window_size: 窗口大小
            
        Returns:
            滤波后的数据
        """
        return signal.medfilt(data, window_size)
    
    def _calculate_dynamic_thresholds(self, energy: np.ndarray) -> Tuple[float, float]:
        """
        计算动态阈值
        
        Args:
            energy: 能量序列
            
        Returns:
            (高阈值, 低阈值)
        """
        # 使用前导无话段估算噪声特性
        noise_length = min(len(energy) // 10, 50)  # 取前10%作为噪声段
        noise_energy = energy[:noise_length]
        
        # 计算噪声的统计特性
        noise_mean = np.mean(noise_energy)
        noise_std = np.std(noise_energy)
        
        # 动态阈值
        energy_threshold_high = noise_mean + self.energy_threshold_high * noise_std
        energy_threshold_low = noise_mean + self.energy_threshold_low * noise_std
        
        return energy_threshold_high, energy_threshold_low
    
    def _detect_boundaries(self, 
                          energy: np.ndarray, 
                          zcr: np.ndarray,
                          energy_threshold_high: float,
                          energy_threshold_low: float) -> Tuple[int, int]:
        """
        检测边界
        
        Args:
            energy: 能量序列
            zcr: 过零率序列
            energy_threshold_high: 能量高阈值
            energy_threshold_low: 能量低阈值
            
        Returns:
            (起始帧, 结束帧)
        """
        num_frames = len(energy)
        
        # 标记有话段
        speech_frames = np.zeros(num_frames, dtype=bool)
        
        for i in range(num_frames):
            # 能量条件
            energy_condition = energy[i] > energy_threshold_low
            # 过零率条件
            zcr_condition = zcr[i] < self.zero_crossing_threshold
            
            speech_frames[i] = energy_condition and zcr_condition
        
        # 寻找起始点
        start_frame = 0
        for i in range(num_frames):
            if speech_frames[i]:
                start_frame = i
                break
        
        # 寻找结束点
        end_frame = num_frames - 1
        for i in range(num_frames - 1, -1, -1):
            if speech_frames[i]:
                end_frame = i
                break
        
        # 应用最小长度约束
        if end_frame - start_frame < self.min_speech_frames:
            # 如果检测到的段太短，扩大搜索范围
            start_frame = max(0, start_frame - self.min_speech_frames // 2)
            end_frame = min(num_frames - 1, end_frame + self.min_speech_frames // 2)
        
        return start_frame, end_frame
    
    def plot_detection_results(self, 
                              signal_data: np.ndarray,
                              start_sample: int,
                              end_sample: int):
        """
        绘制检测结果
        
        Args:
            signal_data: 原始信号
            start_sample: 起始采样点
            end_sample: 结束采样点
        """
        try:
            import matplotlib.pyplot as plt
            
            # 预处理信号
            processed_signal = self._preprocess_signal(signal_data)
            
            # 计算特征
            energy = self._calculate_short_time_energy(processed_signal)
            zcr = self._calculate_zero_crossing_rate(processed_signal)
            
            # 平滑处理
            energy_smooth = self._median_filter(energy)
            zcr_smooth = self._median_filter(zcr)
            
            # 动态阈值
            energy_threshold_high, energy_threshold_low = self._calculate_dynamic_thresholds(energy_smooth)
            
            # 创建时间轴
            time_signal = np.arange(len(signal_data)) / self.sample_rate
            time_frames = np.arange(len(energy)) * self.frame_shift / self.sample_rate
            
            # 绘制结果
            fig, axes = plt.subplots(3, 1, figsize=(12, 10))
            
            # 原始信号
            axes[0].plot(time_signal, signal_data, 'b-', linewidth=1)
            axes[0].axvline(start_sample / self.sample_rate, color='r', linestyle='--', label='起始点')
            axes[0].axvline(end_sample / self.sample_rate, color='r', linestyle='--', label='结束点')
            axes[0].set_title('原始信号')
            axes[0].set_xlabel('时间 (s)')
            axes[0].set_ylabel('振幅')
            axes[0].legend()
            axes[0].grid(True, alpha=0.3)
            
            # 短时能量
            axes[1].plot(time_frames, energy_smooth, 'g-', linewidth=2, label='平滑能量')
            axes[1].axhline(energy_threshold_high, color='r', linestyle='--', label='高阈值')
            axes[1].axhline(energy_threshold_low, color='orange', linestyle='--', label='低阈值')
            axes[1].set_title('短时能量')
            axes[1].set_xlabel('时间 (s)')
            axes[1].set_ylabel('能量')
            axes[1].legend()
            axes[1].grid(True, alpha=0.3)
            
            # 过零率
            axes[2].plot(time_frames, zcr_smooth, 'm-', linewidth=2, label='平滑过零率')
            axes[2].axhline(self.zero_crossing_threshold, color='r', linestyle='--', label='过零率阈值')
            axes[2].set_title('过零率')
            axes[2].set_xlabel('时间 (s)')
            axes[2].set_ylabel('过零率')
            axes[2].legend()
            axes[2].grid(True, alpha=0.3)
            
            plt.tight_layout()
            plt.show()
            
        except ImportError:
            logger.warning("matplotlib未安装，无法绘制检测结果")
    
    def get_detection_parameters(self) -> dict:
        """
        获取检测参数
        
        Returns:
            参数字典
        """
        return {
            'energy_threshold_high': self.energy_threshold_high,
            'energy_threshold_low': self.energy_threshold_low,
            'zero_crossing_threshold': self.zero_crossing_threshold,
            'min_silence_length': self.min_silence_length,
            'min_speech_length': self.min_speech_length,
            'center_clip_threshold': self.center_clip_threshold,
            'frame_length': self.frame_length,
            'frame_shift': self.frame_shift,
            'sample_rate': self.sample_rate
        }
    
    def set_detection_parameters(self, **kwargs):
        """
        设置检测参数
        
        Args:
            **kwargs: 要设置的参数
        """
        for key, value in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, value)
                logger.info(f"设置{key}={value}")
            else:
                logger.warning(f"未知参数: {key}")
        
        # 更新帧数约束
        if 'min_silence_length' in kwargs or 'min_speech_length' in kwargs or 'sample_rate' in kwargs:
            self.min_silence_frames = int(self.min_silence_length * self.sample_rate / self.frame_shift)
            self.min_speech_frames = int(self.min_speech_length * self.sample_rate / self.frame_shift)
    
    def __str__(self) -> str:
        return (f"EndpointDetection(energy_threshold={self.energy_threshold_low}-{self.energy_threshold_high}, "
                f"zcr_threshold={self.zero_crossing_threshold})")
    
    def __repr__(self) -> str:
        return self.__str__() 